package mutex

import (
	"context"
	"sync"
	"sync/atomic"
	"time"
)

var lockValueMap sync.Map

type lockValueData struct {
	// 锁
	mu *sync.Mutex

	// 持有mutex的数量,当持有数量为0时，会清理掉
	holderCnt *atomic.Int64
}

// LockValue locks the value
// ctx 支持超时，取消
func LockValue[T ~int | ~uint | ~int32 | ~uint32 | ~int64 | ~uint16 | ~string](ctx context.Context, key T) error {
	value, _ := lockValueMap.LoadOrStore(key, &lockValueData{
		mu:        &sync.Mutex{},
		holderCnt: &atomic.Int64{},
	})
	lockData, _ := value.(*lockValueData)
	lockData.holderCnt.Add(1)
	// 双循环检查，防止在add前恰巧被删除了：不然会有如下问题
	// 假设有3个协程需要取锁，第一个进入并释放时第二个刚好执行的holderCnt.add前一个。此时map的key被删除，但二个进入持有的ch是第一个被删除的，当有第三个进入时就会存在第二个和第三个互斥不了
	lockValueMap.Store(key, lockData)

	// 自旋取锁或超时
	for {

		// 检查上下文是否已经超时或被取消
		select {
		case <-ctx.Done():
			lockData.holderCnt.Add(-1)
			return ctx.Err()
		default:
			if lockData.mu.TryLock() {
				return nil
			}
		}
		//  没有case准备好，等待一段时间再重试
		time.Sleep(time.Millisecond * 10)
	}
}

// UnlockValue unlocks the value
func UnlockValue[T ~int | ~uint | ~int32 | ~uint32 | ~int64 | ~uint16 | ~string](key T) {
	value, ok := lockValueMap.Load(key)
	if !ok {
		return
	}
	lockData, ok := value.(*lockValueData)
	if !ok {
		return
	}
	// 释放锁
	lockData.holderCnt.Add(-1)
	lockData.mu.Unlock()
	if lockData.holderCnt.Load() <= 0 {
		// 没有持有者，清除map
		lockValueMap.Delete(key)
	}
}
